import argparse
import os
import pickle

import torch
from tqdm import tqdm


def compare(result_dir):
    cnt = 0
    failure = 0
    failure_files = []
    for file in tqdm(os.listdir(result_dir)):
        if not file.endswith('.p'):
            continue

        data = pickle.load(open(os.path.join(result_dir, file), 'rb'))
        if 'cann_result' in data:
            cnt += 1
            pt_result = data['pt_result']
            cann_result = data['cann_result']
            if not check_tensor_equality(pt_result, cann_result):
                failure += 1
                failure_files.append(file)

    with open(os.path.join(result_dir, 'failure_record'), 'w', encoding='utf-8') as f:
        for failure_file in failure_files:
            f.write(failure_file)
            f.write('\n')
        f.write(f'total: {cnt}, fail: {failure}, accuracy: {(cnt - failure) / float(cnt)}')


def check_tensor_equality(a, b):
    constraint = torch.abs(torch.div(torch.minimum(a, b), 1000))
    diff = torch.abs(a - b)
    without_constraint = torch.ge(diff, constraint)
    num_without_constraint = torch.sum(without_constraint)
    threshold = a.numel() / 1000
    return num_without_constraint <= threshold


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', type=str)
    args = parser.parse_args()
    compare(args.output_dir)
